{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 临时脚本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd ../cpp && make clean && make && cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm_book.tvm_ext.libinfo import load_lib\n",
    "\n",
    "_LIB_EXT, _LIB_EXT_NAME = load_lib(name=\"libtvm_ext.so\", search_path=[\"../cpp/outputs/libs\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm_book.tvm_ext.testing._ffi_api as ffi_api"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:20<00:00,  2.00s/it]\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "from tqdm.asyncio import tqdm\n",
    "\n",
    "async def worker(id, is_open=1):\n",
    "    await asyncio.sleep(2)\n",
    "    # if is_open:\n",
    "    #     # async with semaphore:\n",
    "    #         # print(f\"Task {id} acquiring semaphore\")\n",
    "        \n",
    "    #         # print(f\"Task {id} releasing semaphore\")\n",
    "    # else:\n",
    "    #     print(f\"Task {id} acquiring semaphore\")\n",
    "    #     await asyncio.sleep(2)\n",
    "    #     print(f\"Task {id} releasing semaphore\")\n",
    "\n",
    "semaphore = asyncio.Semaphore(20)  # 允许最多两个协程同时运行\n",
    "semaphore\n",
    "tasks = [asyncio.gather(*[asyncio.create_task(worker(i)) for i in range(20)]) for _ in range(10)]\n",
    "async with asyncio.Semaphore(20):\n",
    "    res = await tqdm.gather(*tasks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "worker-0 has slept for 0.27 seconds\n",
      "worker-0 has slept for 0.67 seconds\n",
      "worker-0 has slept for 0.38 seconds\n",
      "worker-0 has slept for 0.70 seconds\n",
      "worker-0 has slept for 0.76 seconds\n",
      "worker-0 has slept for 0.44 seconds\n",
      "worker-0 has slept for 0.52 seconds\n",
      "worker-0 has slept for 0.77 seconds\n",
      "worker-0 has slept for 0.79 seconds\n",
      "worker-0 has slept for 0.81 seconds\n",
      "worker-0 has slept for 0.37 seconds\n",
      "worker-0 has slept for 0.59 seconds\n",
      "worker-0 has slept for 0.93 seconds\n",
      "worker-0 has slept for 0.89 seconds\n",
      "worker-0 has slept for 0.28 seconds\n",
      "worker-0 has slept for 0.94 seconds\n",
      "worker-0 has slept for 0.30 seconds\n",
      "worker-0 has slept for 0.84 seconds\n",
      "worker-0 has slept for 0.36 seconds\n",
      "worker-0 has slept for 0.37 seconds\n",
      "====\n",
      "3 workers slept in parallel for 0.00 seconds\n",
      "total expected sleep time: 11.98 seconds\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "import random\n",
    "import time\n",
    "\n",
    "async def task(val):\n",
    "    return val * 2\n",
    "\n",
    "async def worker(name, queue, res):\n",
    "    while 1:\n",
    "        # 从队列获取一个“工作项”。\n",
    "        sleep_for = await queue.get()\n",
    "\n",
    "        # 休眠 \"sleep_for\" 秒。\n",
    "        res.append(await task(sleep_for))\n",
    "\n",
    "        # 通知队列“工作项”已被处理。\n",
    "        queue.task_done()\n",
    "\n",
    "        print(f'{name} has slept for {sleep_for:.2f} seconds')\n",
    "\n",
    "\n",
    "async def main():\n",
    "    res = []\n",
    "    # 创建一个用于存储我们的“工作项”的队列。\n",
    "    queue = asyncio.Queue()\n",
    "\n",
    "    # 生成随机时段并将它们放入队列。\n",
    "    total_sleep_time = 0\n",
    "    for _ in range(20):\n",
    "        sleep_for = random.uniform(0.05, 1.0)\n",
    "        total_sleep_time += sleep_for\n",
    "        queue.put_nowait(sleep_for)\n",
    "\n",
    "    # 创建三个工作任务来并发地处理队列。\n",
    "    tasks = []\n",
    "    for i in range(3):\n",
    "        task = asyncio.create_task(worker(f'worker-{i}', queue, res))\n",
    "        tasks.append(task)\n",
    "\n",
    "    # 等待直到队列处理完毕。\n",
    "    started_at = time.monotonic()\n",
    "    await queue.join()\n",
    "    total_slept_for = time.monotonic() - started_at\n",
    "\n",
    "    # 取消我们的工作任务。\n",
    "    for task in tasks:\n",
    "        task.cancel()\n",
    "    # 等待直到所有工作任务都被取消。\n",
    "    await asyncio.gather(*tasks, return_exceptions=True)\n",
    "\n",
    "    print('====')\n",
    "    print(f'3 workers slept in parallel for {total_slept_for:.2f} seconds')\n",
    "    print(f'total expected sleep time: {total_sleep_time:.2f} seconds')\n",
    "    return res\n",
    "\n",
    "\n",
    "ts = await main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "48"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "os.cpu_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
